#include "PCFG_v2.h"
// #include "Pthread_OpenMP/Pthread_jin.h"
#include <omp.h>
#include <pthread.h>
#include <algorithm>
using namespace std;

void PriorityQueue::CalProb(PT &pt)
{
    // 计算PriorityQueue里面一个PT的流程如下：
    // 1. 首先需要计算一个PT本身的概率。例如，L6S1的概率为0.15
    // 2. 需要注意的是，Queue里面的PT不是“纯粹的”PT，而是除了最后一个segment以外，全部被value实例化的PT
    // 3. 所以，对于L6S1而言，其在Queue里面的实际PT可能是123456S1，其中“123456”为L6的一个具体value。
    // 4. 这个时候就需要计算123456在L6中出现的概率了。假设123456在所有L6 segment中的概率为0.1，那么123456S1的概率就是0.1*0.15

    // 计算一个PT本身的概率。后续所有具体segment value的概率，直接累乘在这个初始概率值上
    pt.prob = pt.preterm_prob;

    // index: 标注当前segment在PT中的位置
    int index = 0;


    for (int idx : pt.curr_indices)
    {
        // pt.content[index].PrintSeg();
        if (pt.content[index].type == 1)
        {
            // 下面这行代码的意义：
            // pt.content[index]：目前需要计算概率的segment
            // m.FindLetter(seg): 找到一个letter segment在模型中的对应下标
            // m.letters[m.FindLetter(seg)]：一个letter segment在模型中对应的所有统计数据
            // m.letters[m.FindLetter(seg)].ordered_values：一个letter segment在模型中，所有value的总数目
            pt.prob *= m.letters[m.FindLetter(pt.content[index])].ordered_freqs[idx];
            pt.prob /= m.letters[m.FindLetter(pt.content[index])].total_freq;
            // cout << m.letters[m.FindLetter(pt.content[index])].ordered_freqs[idx] << endl;
            // cout << m.letters[m.FindLetter(pt.content[index])].total_freq << endl;
        }
        if (pt.content[index].type == 2)
        {
            pt.prob *= m.digits[m.FindDigit(pt.content[index])].ordered_freqs[idx];
            pt.prob /= m.digits[m.FindDigit(pt.content[index])].total_freq;
            // cout << m.digits[m.FindDigit(pt.content[index])].ordered_freqs[idx] << endl;
            // cout << m.digits[m.FindDigit(pt.content[index])].total_freq << endl;
        }
        if (pt.content[index].type == 3)
        {
            pt.prob *= m.symbols[m.FindSymbol(pt.content[index])].ordered_freqs[idx];
            pt.prob /= m.symbols[m.FindSymbol(pt.content[index])].total_freq;
            // cout << m.symbols[m.FindSymbol(pt.content[index])].ordered_freqs[idx] << endl;
            // cout << m.symbols[m.FindSymbol(pt.content[index])].total_freq << endl;
        }
        index += 1;
    }
    // cout << pt.prob << endl;
}

// void PriorityQueue::init()
// {
//     // cout << m.ordered_pts.size() << endl;
//     // 用所有可能的PT，按概率降序填满整个优先队列
//     for (PT pt : m.ordered_pts)
//     {
//         for (segment seg : pt.content)
//         {
//             if (seg.type == 1)
//             {
//                 // 下面这行代码的意义：
//                 // max_indices用来表示PT中各个segment的可能数目。例如，L6S1中，假设模型统计到了100个L6，那么L6对应的最大下标就是99
//                 // （但由于后面采用了"<"的比较关系，所以其实max_indices[0]=100）
//                 // m.FindLetter(seg): 找到一个letter segment在模型中的对应下标
//                 // m.letters[m.FindLetter(seg)]：一个letter segment在模型中对应的所有统计数据
//                 // m.letters[m.FindLetter(seg)].ordered_values：一个letter segment在模型中，所有value的总数目
//                 pt.max_indices.emplace_back(m.letters[m.FindLetter(seg)].ordered_values.size());
//             }
//             if (seg.type == 2)
//             {
//                 pt.max_indices.emplace_back(m.digits[m.FindDigit(seg)].ordered_values.size());
//             }
//             if (seg.type == 3)
//             {
//                 pt.max_indices.emplace_back(m.symbols[m.FindSymbol(seg)].ordered_values.size());
//             }
//         }
//         //这个PT在所有pt的概率，非具体格式
//         pt.preterm_prob = float(m.preterm_freq[m.FindPT(pt)]) / m.total_preterm;
//         // pt.PrintPT();
//         // cout << " " << m.preterm_freq[m.FindPT(pt)] << " " << m.total_preterm << " " << pt.preterm_prob << endl;

//         // 计算当前pt的概率，这里是具体的PT
//         CalProb(pt);
//         // 将PT放入优先队列
//         priority.emplace_back(pt);
//     }
//     // cout << "priority size:" << priority.size() << endl;
// }

// 并行归并两个有序的 vector
std::vector<PT> merge_vectors(const std::vector<PT> &vec1, const std::vector<PT> &vec2)
{
  if (vec1.empty())
    return vec2;
  if (vec2.empty())
    return vec1;
  
  std::vector<PT> merged(vec1.size() + vec2.size());
  std::merge(vec1.begin(), vec1.end(), vec2.begin(), vec2.end(), merged.begin(),
             [](const PT &a, const PT &b)
             { return a.prob > b.prob; }); // 按概率降序排序
  return merged;
}

// 并行归并排序
std::vector<PT> parallel_merge_sort(std::vector<std::vector<PT>> &local_queues)
{
  // 如果为空或只有一个队列，直接返回
  if (local_queues.empty())
  {
    return std::vector<PT>();
  }
  if (local_queues.size() == 1)
  {
    return local_queues[0];
  }

  // 使用分治法合并队列
  while (local_queues.size() > 1)
  {
    // 创建临时向量存储合并结果
    std::vector<std::vector<PT>> temp_queues((local_queues.size() + 1) / 2);
    int new_size = temp_queues.size();

// 并行合并相邻队列
#pragma omp parallel for schedule(static)
    for (int i = 0; i < new_size; i++)
    {
      if (2 * i + 1 < local_queues.size())
      {
        // 安全地合并两个队列到临时向量
        temp_queues[i] = merge_vectors(local_queues[2 * i], local_queues[2 * i + 1]);
      }
      else
      {
        // 处理奇数情况
        temp_queues[i] = local_queues[2 * i];
      }
    }

    // 将临时向量交换到原向量
    local_queues = std::move(temp_queues);
  }

  // 返回最终合并的队列
  return local_queues.empty() ? std::vector<PT>() : local_queues[0];
}

void PriorityQueue::init()
{
  // 获取线程数
  int num_threads = omp_get_max_threads();
  // int num_threads = 4; // 设置线程数为4

  // 每个线程维护一个本地优先队列
  std::vector<std::vector<PT>> local_queues(num_threads);

// 并行化处理 m.ordered_pts
#pragma omp parallel num_threads(num_threads)
  {
    int thread_id = omp_get_thread_num();
    std::vector<PT> &local_queue = local_queues[thread_id];

// 分块处理 m.ordered_pts
#pragma omp for schedule(static)
    for (size_t i = 0; i < m.ordered_pts.size(); i++)
    {
      PT pt = m.ordered_pts[i];

      // 处理每个 PT 的内容
      for (segment seg : pt.content)
      {
        if (seg.type == 1)
        {
          pt.max_indices.emplace_back(m.letters[m.FindLetter(seg)].ordered_values.size());
        }
        if (seg.type == 2)
        {
          pt.max_indices.emplace_back(m.digits[m.FindDigit(seg)].ordered_values.size());
        }
        if (seg.type == 3)
        {
          pt.max_indices.emplace_back(m.symbols[m.FindSymbol(seg)].ordered_values.size());
        }
      }

      // 计算 PT 的概率
      pt.preterm_prob = float(m.preterm_freq[m.FindPT(pt)]) / m.total_preterm;
      CalProb(pt);

      // 将 PT 添加到本地优先队列
      local_queue.emplace_back(pt);
    }
  }

  // 使用并行归并排序合并所有线程的本地优先队列
  priority = parallel_merge_sort(local_queues);
}

void PriorityQueue::PopNext4()
{
  int NUM_THREADS = 4;
  
  vector<vector<PT>> new_pts(NUM_THREADS);
  vector<string> local_guesses[NUM_THREADS];
  int local_num[NUM_THREADS] = {0};

#pragma omp parallel for schedule(static)
  for (int i = 0; i < NUM_THREADS; i++)
  {
    PT pt;
// 线程安全地获取队首元素
#pragma omp critical
    {
      if (!priority.empty())
      {
        pt = priority.front();
        priority.erase(priority.begin());
      }
    }

    // 如果成功获取到队首元素，则进行处理
    if (!pt.content.empty())
    {
      Generate(pt, local_guesses[i], local_num[i]);
      new_pts[i] = pt.NewPTs();
    }
  }


  //合并猜测结果
  for(int i = 0; i < NUM_THREADS; i++){
    total_guesses += local_num[i];
  }
  guesses.resize(total_guesses);
  for(int i = 0; i < NUM_THREADS; i++){
    guesses.insert(guesses.end(),local_guesses[i].begin(),local_guesses[i].end());
  }

  // 这里我先把各自排成优先队列之后在合并
  //动态分配任务
#pragma omp parallel for schedule(dynamic)
  for (int i = 0; i < NUM_THREADS; i++)
  {
    // 遍历当前线程的 new_pts[i]
    for (PT &pt : new_pts[i])
    {
      // 计算概率
      CalProb(pt);
    }

    // 对当前线程的 new_pts[i] 按概率降序排序
    std::sort(new_pts[i].begin(), new_pts[i].end(), [](const PT &a, const PT &b)
              {
                return a.prob > b.prob; // 按概率降序排序
              });
  }
  
  //先合并
  vector<PT> all_new_pts = parallel_merge_sort(new_pts);
  
  //然后在合并到优先队列中
  priority = merge_vectors(priority, all_new_pts);

      // 然后需要根据即将出队的PT，生成一系列新的PT
      // 这个一些列的猜测本省就是按概率优先队列
  //     for (int i = 0; i < NUM_THREADS; i++)
  // {
  //   // 这个函数的作用是将新的PT插入到优先队列中
  //   // 这里的new_pts[i]是一个vector，存储了当前线程生成的新PT
    
  //   for (PT pt : new_pts[i])
  //   {
  //     // hp计算概率，这里可以分配计算，然后用归并排序
  //     CalProb(pt);
  //     // 接下来的这个循环，作用是根据概率，将新的PT插入到优先队列中
  //     for (auto iter = priority.begin(); iter != priority.end(); iter++)
  //     {
  //       // 对于非队首和队尾的特殊情况
  //       if (iter != priority.end() - 1 && iter != priority.begin())
  //       {
  //         // 判定概率
  //         // 这里判定概率太慢了
  //         if (pt.prob <= iter->prob && pt.prob > (iter + 1)->prob)
  //         {
  //           priority.emplace(iter + 1, pt);
  //           break;
  //         }
  //       }
  //       if (iter == priority.end() - 1)
  //       {
  //         priority.emplace_back(pt);
  //         break;
  //       }
  //       if (iter == priority.begin() && iter->prob < pt.prob)
  //       {
  //         priority.emplace(iter, pt);
  //         break;
  //       }
  //     }
  //   }

    // 现在队首的PT善后工作已经结束，将其出队（删除）
    //万一插到最前面了咋办
    // priority.erase(priority.begin());
  // }
}

// 这个函数你就算看不懂，对并行算法的实现影响也不大
// 当然如果你想做一个基于多优先队列的并行算法，可能得稍微看一看了
vector<PT> PT::NewPTs()
{
    // 存储生成的新PT
    vector<PT> res;

    // 假如这个PT只有一个segment
    // 那么这个segment的所有value在出队前就已经被遍历完毕，并作为猜测输出
    // 因此，所有这个PT可能对应的口令猜测已经遍历完成，无需生成新的PT
    if (content.size() == 1)
    {
        return res;
    }
    else
    {
        // 最初的pivot值。我们将更改位置下标大于等于这个pivot值的segment的值（最后一个segment除外），并且一次只更改一个segment
        // 上面这句话里是不是有没看懂的地方？接着往下看你应该会更明白
        int init_pivot = pivot;

        // 开始遍历所有位置值大于等于init_pivot值的segment
        // 注意i < curr_indices.size() - 1，也就是除去了最后一个segment（这个segment的赋值预留给并行环节）
        for (int i = pivot; i < curr_indices.size() - 1; i += 1)
        {
            // curr_indices: 标记各segment目前的value在模型里对应的下标
            curr_indices[i] += 1;

            // max_indices：标记各segment在模型中一共有多少个value
            if (curr_indices[i] < max_indices[i])
            {
                // 更新pivot值
                pivot = i;
                res.emplace_back(*this);
            }

            // 这个步骤对于你理解pivot的作用、新PT生成的过程而言，至关重要
            curr_indices[i] -= 1;
        }
        pivot = init_pivot;
        return res;
    }

    return res;
}


// 这个函数是PCFG并行化算法的主要载体
// 尽量看懂，然后进行并行实现
void PriorityQueue::Generate(PT pt,vector<string>&local_guesses,int &local_num)
{
    // 计算PT的概率，这里主要是给PT的概率进行初始化
    CalProb(pt);

    // 对于只有一个segment的PT，直接遍历生成其中的所有value即可
    if (pt.content.size() == 1)
    {
        // 指向最后一个segment的指针，这个指针实际指向模型中的统计数据
        segment *a;
        // 在模型中定位到这个segment
        if (pt.content[0].type == 1)
        {
            a = &m.letters[m.FindLetter(pt.content[0])];
        }
        if (pt.content[0].type == 2)
        {
            a = &m.digits[m.FindDigit(pt.content[0])];
        }
        if (pt.content[0].type == 3)
        {
            a = &m.symbols[m.FindSymbol(pt.content[0])];
        }
        
        // Multi-thread TODO：
        // 这个for循环就是你需要进行并行化的主要部分了，特别是在多线程&GPU编程任务中
        // 可以看到，这个循环本质上就是把模型中一个segment的所有value，赋值到PT中，形成一系列新的猜测
        // 这个过程是可以高度并行化的
        for (int i = 0; i < pt.max_indices[0]; i += 1)
        {
            string guess = a->ordered_values[i];
            // cout << guess << endl;
            local_guesses.emplace_back(guess);
            local_num += 1;
        }
        
    
  }
  else
  {
      string guess;
      int seg_idx = 0;
      // 这个for循环的作用：给当前PT的所有segment赋予实际的值（最后一个segment除外）
      // segment值根据curr_indices中对应的值加以确定
      // 这个for循环你看不懂也没太大问题，并行算法不涉及这里的加速
      for (int idx : pt.curr_indices)
      {
        if (pt.content[seg_idx].type == 1)
        {
          guess += m.letters[m.FindLetter(pt.content[seg_idx])].ordered_values[idx];
        }
        if (pt.content[seg_idx].type == 2)
        {
          guess += m.digits[m.FindDigit(pt.content[seg_idx])].ordered_values[idx];
        }
        if (pt.content[seg_idx].type == 3)
        {
          guess += m.symbols[m.FindSymbol(pt.content[seg_idx])].ordered_values[idx];
        }
        seg_idx += 1;
        if (seg_idx == pt.content.size() - 1)
        {
          break;
        }
      }

      // 指向最后一个segment的指针，这个指针实际指向模型中的统计数据
      segment *a;
      if (pt.content[pt.content.size() - 1].type == 1)
      {
        a = &m.letters[m.FindLetter(pt.content[pt.content.size() - 1])];
      }
      if (pt.content[pt.content.size() - 1].type == 2)
      {
        a = &m.digits[m.FindDigit(pt.content[pt.content.size() - 1])];
      }
      if (pt.content[pt.content.size() - 1].type == 3)
      {
        a = &m.symbols[m.FindSymbol(pt.content[pt.content.size() - 1])];
      }

      // Multi-thread TODO：
      // 这个for循环就是你需要进行并行化的主要部分了，特别是在多线程&GPU编程任务中
      // 可以看到，这个循环本质上就是把模型中一个segment的所有value，赋值到PT中，形成一系列新的猜测
      // 这个过程是可以高度并行化的
      for (int i = 0; i < pt.max_indices[pt.content.size() - 1]; i += 1)
      {
        string temp = guess + a->ordered_values[i];
        // cout << temp << endl;
        local_guesses.emplace_back(temp);
        local_num += 1;
      }
      

  }
}